package com.sicau.config;

import com.sicau.rag.MDDocumentReader;
import com.sicau.rag.TokenCountTextSplitter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Configuration;

import java.util.List;

/**
 * 创建一个用于加载向量数据库的VectorStore
 */
@Configuration
@Slf4j
public class VectorStoreConfig {
    private final VectorStore vectorStore;
    private final EmbeddingModel embeddingModel;
    private final MDDocumentReader MDDocumentReader;
    private final TokenCountTextSplitter splitter;

    public VectorStoreConfig(VectorStore vectorStore, EmbeddingModel embeddingModel, MDDocumentReader MDDocumentReader, TokenCountTextSplitter splitter) {
        this.vectorStore = vectorStore;
        this.embeddingModel = embeddingModel;
        this.MDDocumentReader = MDDocumentReader;
        this.splitter = splitter;
    }

    /**
     * 创建一个向量数据库并初始化内容
     * @return
     */
    // @Bean
    public VectorStore vectorStore() {
        List<Document> documents = MDDocumentReader.loadDocuments();
        SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
        vectorStore.add(documents);
        return vectorStore;
    }

    /**
     * 初始化向量数据库
     */
    public void init() {
        log.info("加载初始文档");
        List<Document> documents = MDDocumentReader.loadDocuments();
        documents = splitter.split(documents);
        vectorStore.add(documents);
    }
}
